from langchain import PromptTemplate

instruct_dic = {'command': '指令',
                'context': '上下文',
                'face_to': '受众',
                'intonation': '语调',
                'output_require': '响应输出要求',
                'example': '示例',
                'task': '任务'}

in_out_formatter = '输入：{input}\n输出：{output}'


def get_instruct_prompt(*args, **kwargs):
    prompt = ''
    for k, v in instruct_dic.items():
        if k in kwargs and kwargs[k]:
            prompt += f"### {instruct_dic[k]} ###\n{{{k}}}\n\n"
    # An example prompt with multiple input variables
    multiple_input_prompt = PromptTemplate(
        input_variables=[k for k in instruct_dic.keys() if k in kwargs and kwargs[k]],
        template=prompt
    )
    # print(f'kwargs in get_instruct_prompt is {kwargs.keys()}', flush=True)
    # print(f'multiple_input_prompt varialble is {multiple_input_prompt.input_variables}', flush=True)
    instruct = multiple_input_prompt.format(**kwargs)
    return instruct


def get_task_instruct(prompt_dic):
    example_prompt = '\n\n'.join(
        [in_out_formatter.format(input=example_item['input'], output=example_item['output'])
         for example_item in prompt_dic['example']])
    prompt_dic['example'] = example_prompt
    prompt_dic['task'] = in_out_formatter.format(input=input, output='')
    # print(f'prompt_dic in patent_task is {prompt_dic}', flush=True)
    chat_input = get_instruct_prompt(**prompt_dic)
    # print(chat_input, flush=True)
    return chat_input
